{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用 Kruskal 算法生成最小生成树\n",
    "和 Prim 算法类似，我们也是用了 MST 性质，不过此时不是用 \"军对垒\" 的形式，而是类似 \"分而治之\" 的策略，先形成小的连通分量，在形成一个大的连通分量。\n",
    "\n",
    "算法步骤:\n",
    "1. 将所有的边按权从小到大排序。\n",
    "2. 从第一条边开始，依边权递增的顺序查看每一条边。\n",
    "3. 当查看到第 k 条边 (v,w) 时，如果 v,w 分属不同连通分量，则将他们连接，否则查看 k+1 条边。\n",
    "4. 重复 3 直到只剩下一个连通分量。\n",
    "\n",
    "算法时间复杂度：$O(ElogE)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 邻接矩阵用 NumPy 存储\n",
    "def AdjacencyMatrix(Path):\n",
    "    '''\n",
    "    从 Path 中读取邻接矩阵.\n",
    "    Input:\n",
    "    - Path: 邻接矩阵保存地址.\n",
    "    \n",
    "    Return:\n",
    "    - E: 邻接矩阵.\n",
    "    - V: 顶点集合.\n",
    "    '''\n",
    "    E = None\n",
    "    V = None\n",
    "    index = 0\n",
    "    with open(Path,'r') as f:\n",
    "        while True:\n",
    "            data = f.readline()\n",
    "            if len(data) == 0:\n",
    "                break\n",
    "            # 删除换行符\n",
    "            data = data.rstrip('\\n')\n",
    "            # 分词\n",
    "            data = data.split('\\t')\n",
    "            if index == 0:\n",
    "                E = np.zeros((len(data),len(data)))\n",
    "                E[index] = np.array(data,dtype=int)\n",
    "                V = np.array(list(range(len(data))))\n",
    "            else:\n",
    "                E[index] = np.array(data,dtype=int)\n",
    "            index = index + 1\n",
    "    return E,V\n",
    "    \n",
    "Path = 'example.txt'\n",
    "# 读取邻接矩阵和顶点集\n",
    "E,V = AdjacencyMatrix(Path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def EdgeSort(E):\n",
    "    '''\n",
    "    对边按权值排序. 因为没有好的数据结构，就使用插入排序好了.\n",
    "    Input:\n",
    "    - E: 邻接矩阵.\n",
    "    \n",
    "    Return:\n",
    "    - ind: 坐标序列.\n",
    "    '''\n",
    "    ind = []\n",
    "    for i in range(E.shape[0]):\n",
    "        for j in range(i):\n",
    "            ind.append((i,j))\n",
    "            for k in reversed(range(len(ind))):\n",
    "                if E[ind[k]] < E[ind[k-1]]:\n",
    "                    ind[k],ind[k-1] = ind[k-1],ind[k]\n",
    "                if k == 1:\n",
    "                    break\n",
    "    return ind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def FindComp(C,i):\n",
    "    '''\n",
    "    查找连通分量.\n",
    "    Input:\n",
    "    - C: 连通分量集合.\n",
    "    - i: 顶点.\n",
    "    \n",
    "    Return:\n",
    "    - ind: 连通分量的序号.\n",
    "    '''\n",
    "    for ind in range(len(C)):\n",
    "        for j in range(len(C[ind])):\n",
    "            if C[ind][j] == i:\n",
    "                return ind\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Kruskal(E):\n",
    "    '''\n",
    "    使用 Kruskal 算法生成最小生成树.\n",
    "    Input:\n",
    "    - E: 邻接矩阵\n",
    "    \n",
    "    Return:\n",
    "    - nE: 新的邻接矩阵，代表最小生成树的连接关系，只有 0 和 1.\n",
    "    '''\n",
    "    nE = np.zeros_like(E,dtype=int)\n",
    "    # 对 E 的边按权值排序\n",
    "    indexs = EdgeSort(E)\n",
    "    # 初始化连通分量\n",
    "    C = [[i] for i in range(int(E.shape[0] * (E.shape[0]-1)/2))]\n",
    "    for k,ind in enumerate(indexs):\n",
    "        if len(C)==1:\n",
    "            break\n",
    "        x,y = ind\n",
    "        c_x = FindComp(C,x) # x 的连通分量\n",
    "        c_y = FindComp(C,y) # y 的连通分量\n",
    "        if c_x != c_y:\n",
    "            # 不在同一个连通分量上\n",
    "            C[c_x] = C[c_x] + C[c_y]\n",
    "            C.pop(c_y)\n",
    "            nE[x,y] = nE[y,x] = 1\n",
    "    return nE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0]\n",
      "[0 0]\n",
      "[0 0 0]\n",
      "[1 0 0 0]\n",
      "[0 0 0 0 0]\n",
      "[0 1 0 0 1 0]\n",
      "[0 0 0 0 0 1 0]\n",
      "[0 1 1 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0]\n",
      "[0 0 0 0 0 1 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 1 1 1 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0]\n"
     ]
    }
   ],
   "source": [
    "nE = Kruskal(E)\n",
    "# 以要求格式输出\n",
    "for i in range(E.shape[0]):\n",
    "    print(nE[i,:i+1])\n",
    "\n",
    "# 输出到文件中\n",
    "with open('kruskal_output.txt','w') as f:\n",
    "    for i in range(E.shape[0]):\n",
    "        for j in range(i+1):\n",
    "            if j!=i:\n",
    "                f.write('%s\\t' % nE[i,j])\n",
    "            else:\n",
    "                f.write('%s\\n' % nE[i,j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
